import asyncio
import json
import logging
import os
import pickle
from collections import Counter
from enum import IntEnum
from typing import Iterable, Literal, override

import aiofiles
import httpx
from tqdm import tqdm


def split_range(start, end, chunksize=500):
    chunks = []
    current_start = start
    while current_start < end:
        current_end = min(current_start + chunksize, end)
        chunks.append((current_start, current_end))
        current_start = current_end
    return chunks


class Status(IntEnum):
    PENDING = 0
    COMPLETED = 1
    FAILED = 2
    CANCELLED = 3

class FailedSurpassTolerance(Exception):
    def __init__(self, 
                 error: Exception,
                 n_failed: int = 0, 
                 n_success:int = 0,
                 message="Number of failed tasks surpassed the tolerance"):
        self.error = error
        self.n_failed = n_failed
        self.n_success = n_success
        self.message = f"{message}: {n_failed}"
        super().__init__(self.message)

class AsyncWorker:
    def __init__(self, 
                 name:str,
                 params:list, 
                 client: httpx.AsyncClient, 
                 fp: str,
                 auto_resume:bool = False
                 ):
        self.name = name
        self.params = params
        self.status = [Status.PENDING for _ in params]
        self.client = client
        self.fp = fp
        self._writer = None
        self.checkpoint_path = f'{self.name}.pkl'

        if auto_resume and os.path.exists(self.checkpoint_path):
            self.load_checkpoint()

    def load_checkpoint(self, path:str | None = None):
        if path is None:
            path = self.checkpoint_path
        with open(path, 'rb') as f:
            self.status = pickle.load(f)
        print(f"Read checkpoint file at {path}. ")
        return self
        
    def save_checkpoint(self):
        path = self.checkpoint_path
        with open(path, 'wb') as f:
            pickle.dump(self.status, f)
        print(f"Save checkpoint file at {path}. You can resume task with it.")
        return path


    async def work(self, param):
        raise NotImplementedError


    async def execute_one(self, idx:int):
        if self.status[idx] == Status.COMPLETED:
            return
        param = self.params[idx]
        try:
            data = await self.work(param)
            await self._writer.write(json.dumps(data, ensure_ascii=False) + '\n')
            self.status[idx] = Status.COMPLETED
        except Exception as e:
            self.status[idx] = Status.FAILED
            raise e
        
    
    async def execute_epoch(self, 
                            indexs:Iterable[int], 
                            epoch_desc:str = '',
                            on_errors:Literal['ignore','logging','raise'] = 'ignore',
                            fail_tol_rate:float = 0.1,
                            pbar: tqdm | None = None
                            ):
        for idx in indexs:
            self.status[idx] = Status.PENDING
        
        tasks = [self.execute_one(idx) for idx in indexs]
        if pbar is None:
            pbar = tqdm(total=len(indexs))
        n_success = 0 
        n_failed = 0
        fail_tol = len(indexs) * fail_tol_rate
        async with aiofiles.open(self.fp,'a') as writer:
            self._writer = writer
            with pbar:
                for future in asyncio.as_completed(tasks):
                    try:
                        await future
                        n_success += 1
                    except Exception as e:
                        n_failed += 1
                        if on_errors == 'raise':
                            raise e 
                        if n_failed > fail_tol:
                            raise FailedSurpassTolerance(e, n_failed=n_failed,n_success=n_success)
                        if on_errors == 'logging':
                            logging.exception(e)
                    pbar.update()
                    pbar.set_description(f'{epoch_desc} Success: {n_success}, Failed: {n_failed}')
            self._writer = None
        return n_success, n_failed

    async def run(self, 
                          chunk_size:int = 500, 
                          warmup_size:int = 10,
                          fail_tol_rate:float = 0.1
                          ):
        current_idx = 0
        for idx,status in enumerate(self.status):
            if status != Status.COMPLETED:
                current_idx = idx 
                break
        else:
            return

        try:
            # test one stage
            print(f"Test-One Stage started from: {current_idx}")
            await self.execute_epoch([current_idx], on_errors='raise')
            print("Test-One Stage completed.")

            current_idx += 1

            # test fews stage
            print(f"Test-Fews Stage started from: {current_idx}. Size: {warmup_size}")
            await self.execute_epoch(range(current_idx, current_idx + warmup_size), on_errors='raise')
            print("Test-Fews Stage completed.")

            current_idx += warmup_size

            # main stage
            chunks = split_range(current_idx, len(self.params), chunksize=chunk_size)
            n_chunks = len(chunks)
            print(f"Main Stage started from: {current_idx}. Number of chunks: {n_chunks}, Chunk Size: {chunk_size}")
            with tqdm(total=n_chunks, position=0) as epoch_pbar:
                for i, (start, end) in enumerate(chunks,start=1):
                    try:
                        epoch_pbar.set_description(f'Main Stage: Running "Epoch {i}"')
                        pbar = tqdm(total=end-start, position=1, leave=False)
                        n_success, n_failed = await self.execute_epoch(range(start, end),
                                                                    pbar=pbar,
                                                                    epoch_desc=f'"Epoch {i}" on [{start},{end})',
                                                                    fail_tol_rate=fail_tol_rate)
                        epoch_pbar.update()
                    except FailedSurpassTolerance as e:
                        print(f'"Epoch {i}" Failed too many times! Success: {e.n_success}, Failed: {e.n_failed}.')
                        raise e.error
                    current_idx = end

            counter = Counter(self.status)
            print(f"Main Stage completed. Success: {counter[Status.COMPLETED]}, Failed: {counter[Status.FAILED]}.")

            # retry failed stage
            fails = list(filter(lambda idx:self.status[idx] == Status.FAILED, range(len(self.status))))
            print(f"Retry Stage started. Number of failure: {len(fails)}")
            n_success, n_failed = await self.execute_epoch(fails)
            print(f"\nRetry Stage completed. Success: {n_success}, Failed: {n_failed}.")

            counter = Counter(self.status)
            print(f"All Tasks Completed. Success: {counter[Status.COMPLETED]}, Failed: {counter[Status.FAILED]}.")

        except Exception as E:
            self.save_checkpoint()
            raise E
        

    def run_sync(self, 
            chunk_size:int = 500, 
            warmup_size:int = 10,
            fail_tol_rate:float = 0.1):
        try:
            asyncio.run(self.run(
                chunk_size=chunk_size,
                warmup_size=warmup_size,
                fail_tol_rate=fail_tol_rate
            ))
        except KeyboardInterrupt as e:
            print(repr(e))
            self.save_checkpoint()